package com.springboot.frame.starter.xss;

import lombok.Setter;
import org.springframework.util.AntPathMatcher;
import org.springframework.util.PathMatcher;
import org.springframework.util.StringUtils;
import org.springframework.web.filter.OncePerRequestFilter;
import org.springframework.web.util.HtmlUtils;
import org.springframework.web.util.UrlPathHelper;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

/**
 * @author madman
 */
public class XssFilter extends OncePerRequestFilter {

  private final PathMatcher pathMatcher = new AntPathMatcher();

  private final UrlPathHelper urlPathHelper = new UrlPathHelper();

  @Setter
  private List<String> whiteUrlPatterns = new LinkedList<>();

  @Override
  protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response,
                                  FilterChain filterChain)
      throws ServletException, IOException {
    String lookupPath = urlPathHelper.getLookupPathForRequest(request);
    for (String whiteUrlPattern : whiteUrlPatterns) {
      if (pathMatcher.match(whiteUrlPattern, lookupPath)) {
        filterChain.doFilter(request, response);
        return;
      }
    }
    filterChain.doFilter(new XssHttpServletRequestWrapper(request), response);
  }

  private static class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    private XssHttpServletRequestWrapper(HttpServletRequest servletRequest) {
      super(servletRequest);
    }

    @Override
    public String getHeader(String name) {
      String header = super.getHeader(name);
      if (!StringUtils.isEmpty(header)) {
        header = HtmlUtils.htmlEscape(header);
      }
      return header;
    }

    @Override
    public String getQueryString() {
      String queryString = super.getQueryString();
      if (!StringUtils.isEmpty(queryString)) {
        queryString = HtmlUtils.htmlEscape(queryString);
      }
      return queryString;
    }

    @Override
    public String getParameter(String parameter) {
      String input = super.getParameter(parameter);
      if (!StringUtils.isEmpty(input)) {
        input = HtmlUtils.htmlEscape(input);
      }
      return input;
    }

    @Override
    public String[] getParameterValues(String parameter) {
      String[] values = super.getParameterValues(parameter);
      if (!(values == null || values.length == 0)) {
        for (int i = 0; i < values.length; i++) {
          if (!StringUtils.isEmpty(values[i])) {
            values[i] = HtmlUtils.htmlEscape(values[i]);
          }
        }
      }
      return values;
    }

    @Override
    public Map<String, String[]> getParameterMap() {
      Map<String, String[]> paramMap = new HashMap<>(super.getParameterMap());
      for (Map.Entry<String, String[]> entry : paramMap.entrySet()) {
        String[] values = entry.getValue();
        String[] after = new String[values.length];
        int index = 0;
        for (String value : values) {
          if (!StringUtils.isEmpty(value)) {
            after[index++] = HtmlUtils.htmlEscape(value);
          } else {
            after[index++] = value;
          }
        }
        entry.setValue(after);
      }
      return paramMap;
    }
  }

}
